#include "record.hpp"
#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/time.h>
#include <mutex>

// 定义 MPP_ALIGN 宏
#define MPP_ALIGN(x, a) (((x) + (a) - 1) & ~((a) - 1))
#define FPS 30

static MppEncContext mpp_enc;
static std::mutex record_mutex;
static int64_t last_dts = -1; // 跟踪上一个 DTS

int record_init(uint32_t width, uint32_t height, const char* output_file) {
    std::lock_guard<std::mutex> lock(record_mutex);
    MPP_RET ret = MPP_OK;

    // 重置全局状态
    mpp_enc = {0}; // 清零整个结构体
    last_dts = -1; // 重置 DTS
    mpp_enc.pts = 0; // 重置 PTS
    mpp_enc.ctx = NULL;
    mpp_enc.mpi = NULL;
    mpp_enc.buf_grp = NULL;
    mpp_enc.cfg = NULL;
    mpp_enc.fmt_ctx = NULL;
    mpp_enc.video_stream = NULL;
    mpp_enc.width = width;
    mpp_enc.height = height;
    mpp_enc.hor_stride = MPP_ALIGN(width, 16);
    mpp_enc.ver_stride = MPP_ALIGN(height, 16);

    // 初始化 FFmpeg
    avformat_alloc_output_context2(&mpp_enc.fmt_ctx, NULL, "mp4", output_file);
    if (!mpp_enc.fmt_ctx) {
        printf("无法分配 FFmpeg 输出上下文\n");
        return -1;
    }

    // 创建视频流
    const AVOutputFormat *fmt = mpp_enc.fmt_ctx->oformat;
    mpp_enc.video_stream = avformat_new_stream(mpp_enc.fmt_ctx, NULL);
    if (!mpp_enc.video_stream) {
        printf("无法创建视频流\n");
        avformat_free_context(mpp_enc.fmt_ctx);
        return -1;
    }

    // 设置视频流参数
    const AVCodec *codec = avcodec_find_encoder(AV_CODEC_ID_H264);
    if (!codec) {
        printf("找不到 H.264 编码器\n");
        avformat_free_context(mpp_enc.fmt_ctx);
        return -1;
    }
    mpp_enc.video_stream->codecpar->codec_id = AV_CODEC_ID_H264;
    mpp_enc.video_stream->codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
    mpp_enc.video_stream->codecpar->width = width;
    mpp_enc.video_stream->codecpar->height = height;
    mpp_enc.video_stream->codecpar->format = AV_PIX_FMT_YUV420P;
    mpp_enc.video_stream->time_base = (AVRational){1, FPS};
    mpp_enc.video_stream->codecpar->bit_rate = 10000000; // 10 Mbps，降低比特率

    // 打开输出文件
    if (!(fmt->flags & AVFMT_NOFILE)) {
        if (avio_open(&mpp_enc.fmt_ctx->pb, output_file, AVIO_FLAG_WRITE) < 0) {
            printf("无法打开输出文件\n");
            avformat_free_context(mpp_enc.fmt_ctx);
            return -1;
        }
    }

    // 写入 MP4 头部
    if (avformat_write_header(mpp_enc.fmt_ctx, NULL) < 0) {
        printf("无法写入 MP4 头部\n");
        avformat_free_context(mpp_enc.fmt_ctx);
        return -1;
    }

    // 初始化 MPP
    mpp_enc.width = width;
    mpp_enc.height = height;
    mpp_enc.hor_stride = MPP_ALIGN(width, 16);
    mpp_enc.ver_stride = MPP_ALIGN(height, 16);

    ret = mpp_create(&mpp_enc.ctx, &mpp_enc.mpi);
    if (ret != MPP_OK) {
        printf("mpp_create 失败\n");
        avformat_free_context(mpp_enc.fmt_ctx);
        return -1;
    }

    ret = mpp_init(mpp_enc.ctx, MPP_CTX_ENC, MPP_VIDEO_CodingAVC);
    if (ret != MPP_OK) {
        printf("mpp_init 失败\n");
        mpp_destroy(mpp_enc.ctx);
        avformat_free_context(mpp_enc.fmt_ctx);
        return -1;
    }

    mpp_enc_cfg_init(&mpp_enc.cfg);
    mpp_enc.mpi->control(mpp_enc.ctx, MPP_ENC_GET_CFG, mpp_enc.cfg);

    mpp_enc_cfg_set_s32(mpp_enc.cfg, "prep:width", mpp_enc.width);
    mpp_enc_cfg_set_s32(mpp_enc.cfg, "prep:height", mpp_enc.height);
    mpp_enc_cfg_set_s32(mpp_enc.cfg, "prep:hor_stride", mpp_enc.hor_stride);
    mpp_enc_cfg_set_s32(mpp_enc.cfg, "prep:ver_stride", mpp_enc.ver_stride);
    mpp_enc_cfg_set_s32(mpp_enc.cfg, "prep:format", MPP_FMT_YUV420SP);
    mpp_enc_cfg_set_s32(mpp_enc.cfg, "rc:mode", MPP_ENC_RC_MODE_CBR);
    mpp_enc_cfg_set_s32(mpp_enc.cfg, "rc:fps_in_denorm", 1);
    mpp_enc_cfg_set_s32(mpp_enc.cfg, "rc:fps_out_denorm", 1);
    mpp_enc_cfg_set_s32(mpp_enc.cfg, "rc:fps_in_num", FPS);
    mpp_enc_cfg_set_s32(mpp_enc.cfg, "rc:fps_out_num", FPS);
    mpp_enc_cfg_set_s32(mpp_enc.cfg, "rc:bps_target", 10000000); // 10 Mbps
    mpp_enc_cfg_set_s32(mpp_enc.cfg, "codec:type", MPP_VIDEO_CodingAVC);
    mpp_enc_cfg_set_s32(mpp_enc.cfg, "h264:profile", 100);
    mpp_enc_cfg_set_s32(mpp_enc.cfg, "h264:level", 40);
    mpp_enc_cfg_set_s32(mpp_enc.cfg, "h264:cabac_en", 1);
    mpp_enc_cfg_set_s32(mpp_enc.cfg, "h264:cabac_idc", 0);
    mpp_enc_cfg_set_s32(mpp_enc.cfg, "h264:trans8x8", 1);
    mpp_enc_cfg_set_s32(mpp_enc.cfg, "rc:gop", FPS / 2); // 每 0.5 秒 I 帧
    mpp_enc_cfg_set_s32(mpp_enc.cfg, "rc:low_delay", 1); // 低延迟模式

    mpp_enc.mpi->control(mpp_enc.ctx, MPP_ENC_SET_CFG, mpp_enc.cfg);

    // 获取 H.264 头部（SPS/PPS）并写入 MP4
    MppPacket header_packet = NULL;
    mpp_enc.mpi->control(mpp_enc.ctx, MPP_ENC_GET_EXTRA_INFO, &header_packet);
    if (header_packet) {
        AVPacket *avpkt = av_packet_alloc();
        if (!avpkt) {
            printf("无法为头部分配 AVPacket\n");
            mpp_packet_deinit(&header_packet);
            avformat_free_context(mpp_enc.fmt_ctx);
            return -1;
        }
        avpkt->data = (uint8_t*)mpp_packet_get_data(header_packet);
        avpkt->size = mpp_packet_get_length(header_packet);
        avpkt->pts = AV_NOPTS_VALUE;
        avpkt->dts = AV_NOPTS_VALUE;
        avpkt->stream_index = mpp_enc.video_stream->index;
        av_write_frame(mpp_enc.fmt_ctx, avpkt);
        av_packet_unref(avpkt);
        av_packet_free(&avpkt);
        mpp_packet_deinit(&header_packet);
    }

    ret = mpp_buffer_group_get_internal(&mpp_enc.buf_grp, MPP_BUFFER_TYPE_ION);
    if (ret != MPP_OK) {
        printf("mpp_buffer_group_get 失败\n");
        avformat_free_context(mpp_enc.fmt_ctx);
        return -1;
    }

    mpp_enc.pts = 0;
    return 0;
}

void record_close() {
    std::lock_guard<std::mutex> lock(record_mutex);

    // 刷新编码器
    if (mpp_enc.ctx && mpp_enc.mpi) {
        MppFrame frame = NULL;
        MPP_RET ret = mpp_frame_init(&frame);
        if (ret == MPP_OK) {
            mpp_frame_set_eos(frame, 1);
            mpp_enc.mpi->encode_put_frame(mpp_enc.ctx, frame);
            mpp_frame_deinit(&frame);
        }

        while (1) {
            MppPacket packet = NULL;
            MPP_RET ret = mpp_enc.mpi->encode_get_packet(mpp_enc.ctx, &packet);
            if (ret != MPP_OK || !packet) {
                if (ret != MPP_OK) {
                    printf("encode_get_packet 失败：%d\n", ret);
                }
                break;
            }

            AVPacket *avpkt = av_packet_alloc();
            if (!avpkt) {
                printf("无法分配 AVPacket\n");
                mpp_packet_deinit(&packet);
                continue;
            }

            avpkt->data = (uint8_t*)mpp_packet_get_data(packet);
            avpkt->size = mpp_packet_get_length(packet);
            avpkt->stream_index = mpp_enc.video_stream->index;

            // 获取 PTS
            int64_t packet_pts = mpp_packet_get_pts(packet);
            if (packet_pts == -1 || packet_pts < 0) {
                packet_pts = mpp_enc.pts;
            }
            avpkt->pts = av_rescale_q(packet_pts, (AVRational){1, 1000000}, mpp_enc.video_stream->time_base);

            // 确保 DTS 单调递增
            int64_t scaled_dts = avpkt->pts;
            if (scaled_dts <= last_dts) {
                scaled_dts = last_dts + 1;
            }
            avpkt->dts = scaled_dts;
            last_dts = scaled_dts;

            // ret = (MPP_RET)av_write_frame(mpp_enc.fmt_ctx, avpkt);
            int av_ret = av_write_frame(mpp_enc.fmt_ctx, avpkt);
            if (av_ret < 0) {
/**************************************不知道为啥这里会报错************************************************************** */
                // printf("写入帧失败\n");
                av_packet_unref(avpkt);
                av_packet_free(&avpkt);
                mpp_packet_deinit(&packet);
                continue;
            }

            av_packet_unref(avpkt);
            av_packet_free(&avpkt);
            mpp_packet_deinit(&packet);
            mpp_enc.pts += 1000000 / FPS;
        }
    }

    // 写入 MP4 尾部
    if (mpp_enc.fmt_ctx) {
        av_write_trailer(mpp_enc.fmt_ctx);
    }

    // 释放 FFmpeg 资源
    if (mpp_enc.fmt_ctx && !(mpp_enc.fmt_ctx->oformat->flags & AVFMT_NOFILE)) {
        avio_closep(&mpp_enc.fmt_ctx->pb);
    }
    if (mpp_enc.fmt_ctx) {
        avformat_free_context(mpp_enc.fmt_ctx);
        mpp_enc.fmt_ctx = NULL;
    }

    // 释放 MPP 资源
    if (mpp_enc.ctx) {
        mpp_destroy(mpp_enc.ctx);
        mpp_enc.ctx = NULL;
    }
    if (mpp_enc.buf_grp) {
        mpp_buffer_group_put(mpp_enc.buf_grp);
        mpp_enc.buf_grp = NULL;
    }
    if (mpp_enc.cfg) {
        mpp_enc_cfg_deinit(mpp_enc.cfg);
        mpp_enc.cfg = NULL;
    }
}

int record_frame(const cv::Mat& frame) {
    std::lock_guard<std::mutex> lock(record_mutex);
    MPP_RET ret = MPP_OK;
    MppFrame mpp_frame = NULL;
    MppPacket packet = NULL;
    MppBuffer buffer = NULL;

    // 输入验证
    if (frame.empty() || frame.cols != mpp_enc.width || frame.rows != mpp_enc.height || frame.type() != CV_8UC3) {
        printf("无效帧输入：大小=%dx%d，类型=%d\n", frame.cols, frame.rows, frame.type());
        return -1;
    }

    // 将BGR转换为YUV420 (I420)
    cv::Mat yuv_i420;
    cv::cvtColor(frame, yuv_i420, cv::COLOR_RGB2YUV_I420);

    // 计算YUV各分量大小（考虑MPP对齐要求）
    size_t y_size = mpp_enc.hor_stride * mpp_enc.ver_stride;
    size_t uv_stride = mpp_enc.hor_stride / 2;  // 色度分量步长为亮度的一半
    size_t uv_height = mpp_enc.ver_stride / 2;  // 色度高度为亮度的一半
    size_t uv_size = uv_stride * uv_height * 2; // UV总大小
    
    // 分配MPP缓冲区（已考虑对齐）
    ret = mpp_buffer_get(mpp_enc.buf_grp, &buffer, y_size + uv_size);
    if (ret != MPP_OK) {
        printf("mpp_buffer_get 失败：%d\n", ret);
        return -1;
    }

    uint8_t *yuv_data = (uint8_t *)mpp_buffer_get_ptr(buffer);
    uint8_t *y_plane = yuv_data;
    uint8_t *uv_plane = yuv_data + y_size;

    // 计算实际数据量（非对齐大小）
    size_t actual_y_size = mpp_enc.width * mpp_enc.height;
    size_t actual_uv_size = actual_y_size / 4;
    
    /*
     * 内存复制优化：
     * 1. Y平面直接整块复制
     * 2. UV平面需逐行复制考虑步长对齐
     */
    // 复制Y分量
    memcpy(y_plane, yuv_i420.data, actual_y_size);

    // 复制UV分量（考虑步长对齐）
    const uint8_t* u_plane = yuv_i420.data + actual_y_size;
    const uint8_t* v_plane = u_plane + actual_uv_size;
    
    for (size_t row = 0; row < uv_height; row++) {
        // 仅复制有效数据部分，避免越界
        if (row < mpp_enc.height/2) {
            size_t src_offset = row * (mpp_enc.width/2);
            size_t dst_offset = row * uv_stride * 2;
            
            for (size_t col = 0; col < mpp_enc.width/2; col++) {
                uv_plane[dst_offset + 2*col] = u_plane[src_offset + col];   // U分量
                uv_plane[dst_offset + 2*col + 1] = v_plane[src_offset + col]; // V分量
            }
        }
    }

    // 初始化并填充MPP帧
    ret = mpp_frame_init(&mpp_frame);
    if (ret != MPP_OK) {
        printf("mpp_frame_init 失败：%d\n", ret);
        mpp_buffer_put(buffer);
        return -1;
    }

    mpp_frame_set_width(mpp_frame, mpp_enc.width);
    mpp_frame_set_height(mpp_frame, mpp_enc.height);
    mpp_frame_set_hor_stride(mpp_frame, mpp_enc.hor_stride);
    mpp_frame_set_ver_stride(mpp_frame, mpp_enc.ver_stride);
    mpp_frame_set_fmt(mpp_frame, MPP_FMT_YUV420SP);
    mpp_frame_set_pts(mpp_frame, mpp_enc.pts);
    mpp_frame_set_buffer(mpp_frame, buffer);

    // 提交帧到编码器
    ret = mpp_enc.mpi->encode_put_frame(mpp_enc.ctx, mpp_frame);
    if (ret != MPP_OK) {
        printf("encode_put_frame 失败：%d\n", ret);
        mpp_frame_deinit(&mpp_frame);
        mpp_buffer_put(buffer);
        return -1;
    }

    // 处理编码后的数据包
    while (1) {
        ret = mpp_enc.mpi->encode_get_packet(mpp_enc.ctx, &packet);
        if (ret != MPP_OK || !packet) {
            if (ret != MPP_OK) {
                printf("encode_get_packet 失败：%d\n", ret);
            }
            break;
        }

        AVPacket *avpkt = av_packet_alloc();
        if (!avpkt) {
            printf("无法分配AVPacket\n");
            mpp_packet_deinit(&packet);
            continue;
        }

        avpkt->data = (uint8_t*)mpp_packet_get_data(packet);
        avpkt->size = mpp_packet_get_length(packet);
        avpkt->stream_index = mpp_enc.video_stream->index;

        // 处理时间戳
        int64_t packet_pts = mpp_packet_get_pts(packet);
        if (packet_pts == -1) {
            packet_pts = mpp_enc.pts;
        }
        avpkt->pts = av_rescale_q(packet_pts, (AVRational){1, 1000000}, mpp_enc.video_stream->time_base);

        // 确保DTS单调递增
        int64_t scaled_dts = avpkt->pts;
        if (scaled_dts <= last_dts) {
            scaled_dts = last_dts + 1;
        }
        avpkt->dts = scaled_dts;
        last_dts = scaled_dts;

        // 写入封装容器
        av_write_frame(mpp_enc.fmt_ctx, avpkt);
        av_packet_unref(avpkt);
        av_packet_free(&avpkt);
        mpp_packet_deinit(&packet);
    }

    // 更新状态
    mpp_enc.pts += 1000000 / FPS;  // 下一帧PTS
    mpp_frame_deinit(&mpp_frame);   // 释放帧资源
    mpp_buffer_put(buffer);         // 释放缓冲区
    
    return 0;
}

